#-*- coding: utf8 -*-
from timeseries.base_types import TimeSeries
import urllib
import numpy as np
import datetime
import time
import csv

def get_quotes(quote, date_range=None):
    """
    Retrieve historical data from the given 'quote' from Yahoo Finance. 
    (Requires a connection with the Internet.)
    
    Arguments
    ---------
    
    quote: str
        The name of the quote to be downloaded (e.g., 'GOOG', 'YHOO', etc)
    
    date_range: (initial_date, final_date)
        If no date_range is given, it download the full time series. The 
        initial_date and final_date can be given as a datetime.date object
        or a string in the 'YYYY-MM-DD' format.
        
    Output
    ------
    
    A TimeSeries() object.
    """

    # Open the url corresponding to the CSV file for given quote
    if date_range is not None:
        date_range = list(date_range)
        for idx, date in enumerate(date_range):
            try:
                date = date.split('-')
            except AttributeError:
                pass
            else:
                date = datetime.date(*map(int, date))
                date_range[idx] = date

        date_i, date_f = date_range
        url = ('http://ichart.yahoo.com/table.csv?s={quote}&'
               'a={month_i}&b={day_i}&c={year_i}&'
               'd={month_f}&e={day_f}&f={year_f}&'
               'ignore=.csv')
        url = url.format(quote=quote,
            day_i=date_i.day, month_i=date_i.month - 1, year_i=date_i.year,
            day_f=date_f.day, month_f=date_f.month - 1, year_f=date_f.year)
    else:
        url = 'http://ichart.yahoo.com/table.csv?s={quote}'.format(quote=quote)

    try:
        url = urllib.urlopen(url)
        lines = url.readlines()
        #TODO: check if line is not an 404 error
    finally:
        url.close()

    # Read header and save data
    csv_file = csv.reader(lines)
    header = csv_file.next()
    header = [ name.lower() for name in header[1:] ]
    csv_file = list(csv_file)[::-1]

    # Save times
    strptime = datetime.datetime.strptime
    times = (strptime(l[0], '%Y-%m-%d') for l in csv_file)
    times = [ time.mktime(dt.timetuple()) for dt in times ]

    # Convert data to numpy array
    data = [ map(float, l[1:]) for l in csv_file ]
    data = np.array(data, dtype=float)

    return TimeSeries(times, data, names=header)

if __name__ == '__main__':
    from timeseries.visualization import ts_plot
    from pylab import show
    t = get_quotes('GOOG')
    ts_plot(t, idx=['close', 'open']); show()
